import numpy as np
import math
import torch
from torch.utils.data import Dataset
import pandas as pd


def standardization(data):
    mu = np.mean(data, axis=0)
    sigma = np.std(data, axis=0)
    return (data - mu) / sigma


def dataPreprocess(data):
    new_data = []
    nan_index = []
    sum = 0
    for i in np.arange(np.size(data, 0)):
        if math.isnan(data[i]):
            nan_index.append(i)
        else:
            sum += data[i]
    mu = sum / (len(data) - len(nan_index))
    for i in nan_index:
        data[i] = mu
    return data


def dataLoader(data):
    res = []
    for feature in data:
        new_data = dataPreprocess(feature)
        new_data = standardization(new_data)
        res.append(np.asarray(new_data))
    res = np.asarray(res)
    return torch.from_numpy(res.astype(np.float32))


class HousePriceDataSet(Dataset):
    def __init__(self, root, features_name, output_name):
        self.data = pd.read_csv(root)
        self.features = np.asarray([self.data[feature] for feature in features_name])
        self.output_temp = np.asarray(self.data[output_name])
        self.output = np.expand_dims(self.output_temp, axis=0)
        self.features_tensor = dataLoader(self.features).T
        self.output_tensor = dataLoader(self.output).T

    def __len__(self):
        return len(self.output)

    def __getitem__(self, item):
        return self.features_tensor[item], self.output_tensor[item]

